package middleware

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"github.com/golang-jwt/jwt/v4"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
)

type JwtAuthMiddleware struct {
	Secret string
}

func NewJwtAuthMiddleware(secret string) *JwtAuthMiddleware {
	return &JwtAuthMiddleware{
		Secret: secret,
	}
}

func (m *JwtAuthMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 从Authorization header获取token
		authHeader := r.Header.Get("Authorization")
		if authHeader == "" {
			logx.Error("Authorization header is missing")
			httpx.Error(w, fmt.Errorf("未授权：缺少Authorization头"))
			return
		}

		// 检查Bearer前缀
		if !strings.HasPrefix(authHeader, "Bearer ") {
			logx.Error("Authorization header format is invalid")
			httpx.Error(w, fmt.Errorf("未授权：Authorization头格式错误"))
			return
		}

		// 提取token
		tokenString := authHeader[7:] // 去掉"Bearer "前缀
		if tokenString == "" {
			logx.Error("Token is empty")
			httpx.Error(w, fmt.Errorf("未授权：token为空"))
			return
		}

		// 解析token
		claims := jwt.MapClaims{}
		token, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
			// 验证签名方法
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
			}
			return []byte(m.Secret), nil
		})

		if err != nil {
			logx.Errorf("Token parsing failed: %v", err)
			httpx.Error(w, fmt.Errorf("未授权：token解析失败"))
			return
		}

		if !token.Valid {
			logx.Error("Token is invalid")
			httpx.Error(w, fmt.Errorf("未授权：token无效"))
			return
		}

		// 将claims添加到context中
		ctx := context.WithValue(r.Context(), "claims", claims)
		
		// 同时也添加用户ID到context中，方便直接获取
		if userIdClaim, exists := claims["id"]; exists {
			ctx = context.WithValue(ctx, "userId", userIdClaim)
		}

		// 添加用户名到context中
		if usernameClaim, exists := claims["username"]; exists {
			ctx = context.WithValue(ctx, "username", usernameClaim)
		}

		// 添加角色到context中
		if roleClaim, exists := claims["role"]; exists {
			ctx = context.WithValue(ctx, "role", roleClaim)
		}

		logx.Infof("JWT authentication successful for user: %v", claims["username"])

		// 使用新的context继续处理请求
		next(w, r.WithContext(ctx))
	})
}